{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Importing all the required packages\n",
    "import torch as tr\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.io as sio\n",
    "import math\n",
    "import time\n",
    "import sys\n",
    "from matplotlib import animation, rc\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Defining some of the functions used\n",
    "def de2bi(d, n):\n",
    "\td = np.array(d)\n",
    "\tpower = 2**np.arange(n)\n",
    "\treturn (np.floor((d[:,None]%(2*power))/power))\n",
    "\n",
    "def bi2de(L,m):\n",
    "    Ldec = np.zeros((int(2**m),1), dtype = np.int)\n",
    "    for i in range(int(m)):\n",
    "        for j in range(int(2**m)):\n",
    "            Ldec[j] = Ldec[j] + L[j,i]*2**i \n",
    "    return Ldec\n",
    "    \n",
    "def get_labeling(m):\n",
    "    M  = 2**m\n",
    "    if m == 1:\n",
    "        L = np.asarray([[0],[1]])\n",
    "    else:\n",
    "        L = np.zeros((M,m))\n",
    "        L[0:int(M/2),1:m] = get_labeling(m-1)\n",
    "        L[int(M/2):M, 1:m] = np.flipud(L[0:int(M/2),1:m])\n",
    "        L[int(M/2):M,0] = 1\n",
    "    return L\n",
    "\n",
    "def get_constellation(M):\n",
    "    Delta = np.sqrt(3/(2*(M-1)))\n",
    "    Xpam = np.expand_dims(Delta*np.arange(-(np.sqrt(M)-1),np.sqrt(M)+1,2),axis = 1)\n",
    "    xpamt_2D = np.tile(Xpam,(1,int(np.sqrt(M))))\n",
    "    xpamt = np.expand_dims(xpamt_2D.flatten(),axis = 1)\n",
    "    X = np.transpose(np.reshape(np.asarray([xpamt, np.tile(Xpam,(int(np.sqrt(M)),1))]),(2,M)))\n",
    "    Ltmp = get_labeling(int(np.log2(M)/2))\n",
    "    Ltmp_dec = bi2de(Ltmp,int(np.log2(M))/2)\n",
    "    Ltmp_dec2 = np.tile(Ltmp_dec,(1,int(np.sqrt(M))))\n",
    "    Ltmp_dec3 = np.expand_dims(Ltmp_dec2.flatten(),axis = 1)\n",
    "    L = np.concatenate((np.reshape(de2bi(Ltmp_dec3,int(np.log2(M)/2)),(M,int(np.log2(M)/2))), np.tile(Ltmp,(int(np.sqrt(M)),1))), axis = 1)\n",
    "    Ldec = np.reshape(np.asarray(bi2de(np.fliplr(L),int(np.log2(M))),dtype = np.int),M)\n",
    "    return [X,Ldec]\n",
    "\n",
    "def get_APSK(M2,Amount_of_Rings):\n",
    "    M = int(M2/Amount_of_Rings) * np.ones(Amount_of_Rings, dtype = np.int)\n",
    "    X = np.zeros((sum(M),2))\n",
    "    if Amount_of_Rings > 1:\n",
    "        idx = 0\n",
    "        l_r1 = get_labeling(int(np.log2(M[0])))*M.shape[0]\n",
    "        ldec_r1 = bi2de(l_r1,int(np.log2(M[0])))\n",
    "        print('a')\n",
    "        l_rs = get_labeling(int(np.log2(M.shape[0])))\n",
    "        ldec_rs = bi2de(l_rs,int(np.log2(M.shape[0])))\n",
    "        l_apsk = np.zeros((sum(M),1),dtype = np.int)\n",
    "        for i in range(M.shape[0]):\n",
    "            R = np.sqrt(-np.log(1-(i+0.5)*1/M.shape[0]))\n",
    "            for j in range(M[i]):\n",
    "                X[idx+j,:] = [R*np.cos(j*2*math.pi/M[i]), R*np.sin(j*2*math.pi/M[i]) ]\n",
    "            l_apsk[idx:idx+M[i],] = ldec_r1 + ldec_rs[i]\n",
    "            idx = idx + M[i]\n",
    "        l_apsk = np.squeeze(l_apsk)\n",
    "    else:\n",
    "        Lbin = get_labeling(int(np.log2(M[0])))\n",
    "        l_apsk = np.squeeze(bi2de(Lbin,int(np.log2(M[0]))))\n",
    "        for j in range(M2):\n",
    "             X[j,:] = [np.cos(j*2*math.pi/M2), np.sin(j*2*math.pi/M2)]\n",
    "    return [X, l_apsk]\n",
    "\n",
    "def get_constellation_4D(M):\n",
    "    m = int(np.log2(M))\n",
    "    m_pd = int(m/4)\n",
    "    Delta = np.sqrt(np.sqrt(3/(2*(M-1))))\n",
    "    Xpam = np.expand_dims(Delta*np.arange(-(np.sqrt(np.sqrt(M))-1),np.sqrt(np.sqrt(M))+1,2),axis = 1)\n",
    "    Lpam = get_labeling(m_pd)\n",
    "    X1 =  (np.expand_dims(np.ones(int(M/2**m_pd)),axis= 1)*np.transpose(Xpam)).flatten('F')\n",
    "    X2 =  np.tile((np.expand_dims(np.ones(int(M/2**(m_pd*2))),axis= 1)*np.transpose(Xpam)).flatten('F'), 2**m_pd)\n",
    "    X3 =  np.tile((np.expand_dims(np.ones(int(M/2**(m_pd*3))),axis= 1)*np.transpose(Xpam)).flatten('F'), 2**(m_pd*2))\n",
    "    X4 =  np.tile((np.expand_dims(np.ones(int(M/2**(m_pd*4))),axis= 1)*np.transpose(Xpam)).flatten('F'), 2**(m_pd*3))\n",
    "    L1 =  np.reshape(np.tile(Lpam,(int(M/2**m_pd))), (M,m_pd))\n",
    "    L2 =  np.tile(np.reshape(np.tile(Lpam,(int(M/2**(m_pd*2)))), (int(M/2**m_pd),m_pd)),(2**m_pd,1))\n",
    "    L3 =  np.tile(np.reshape(np.tile(Lpam,(int(M/2**(m_pd*3)))), (int(M/2**(m_pd*2)),m_pd)),(2**(m_pd*2),1))\n",
    "    L4 =  np.tile(np.reshape(np.tile(Lpam,(int(M/2**(m_pd*4)))), (int(M/2**(m_pd*3)),m_pd)),(2**(m_pd*3),1))\n",
    "    X = np.transpose(np.asarray([X1,X2,X3,X4]))\n",
    "    Lbin =np.asarray(np.concatenate((L1,L2,L3,L4),axis = 1), dtype = np.int)\n",
    "    Ldec = np.squeeze(bi2de(Lbin,m))\n",
    "    return [X,Ldec]  \n",
    "\n",
    "def get_APSK_4D(M, Amount_of_Rings):\n",
    "    m = int(np.log2(M))\n",
    "    M2 = int(np.sqrt(M))\n",
    "    X_2D, l_apsk_2D = get_APSK(M2,Amount_of_Rings)\n",
    "    Lbin_2D = de2bi(l_apsk_2D, int(m/2))\n",
    "    X1 = np.repeat(X_2D,M2,axis = 0)\n",
    "    X2 = np.reshape(np.repeat(X_2D,M2,axis = 1), (M,2), 'F')\n",
    "    L1 = np.repeat(Lbin_2D, M2, axis = 0)\n",
    "    L2 = np.reshape(np.repeat(Lbin_2D,M2,axis = 1), (M,int(m/2)), 'F')\n",
    "    X = np.concatenate([X1,X2], axis = 1)\n",
    "    Lbin = np.concatenate([L1,L2], axis = 1)\n",
    "    Ldec = np.squeeze(bi2de(Lbin,m))\n",
    "    return[X,Ldec]\n",
    "\n",
    "def get_Optimized_4D(X_2D, l_2D, M):\n",
    "    m = int(np.log2(M)) \n",
    "    M2 = int(np.sqrt(M))\n",
    "    Lbin_2D = de2bi(l_2D, int(m/2))\n",
    "    X1 = np.repeat(X_2D,M2,axis = 0)\n",
    "    X2 = np.reshape(np.repeat(X_2D,M2,axis = 1), (M,2), 'F')\n",
    "    L1 = np.repeat(Lbin_2D, M2, axis = 0)\n",
    "    L2 = np.reshape(np.repeat(Lbin_2D,M2,axis = 1), (M,int(m/2)), 'F')\n",
    "    X = np.concatenate([X1,X2], axis = 1)\n",
    "    Lbin = np.concatenate([L1,L2], axis = 1)\n",
    "    Ldec = np.squeeze(bi2de(Lbin,m))\n",
    "    return[X,Ldec]\n",
    "\n",
    "def get_Random_4D(M):\n",
    "    m = int(np.log2(M))\n",
    "    M2 = int(M/16)\n",
    "    X_G = abs(np.random.normal(0,1,(M2,4)))\n",
    "    L_2D = get_labeling(m)\n",
    "    L_Q = get_labeling(4)\n",
    "    X_T = np.tile(X_G,[16,1])\n",
    "    Q_M = 2*L_Q -1\n",
    "    Q_M = np.repeat(Q_M,M2, axis = 0)\n",
    "    X_RN = X_T*Q_M\n",
    "    L_2D = np.tile(L_2D,[16,1])\n",
    "    L_Q = np.repeat(L_Q,M2,axis = 0)\n",
    "    Lbin = np.concatenate([L_2D,L_Q],axis = 1)\n",
    "    Ldec = np.squeeze(bi2de(Lbin,m))\n",
    "    return[X_RN,Ldec]\n",
    "\n",
    "def awgn_channel(x):\n",
    "    return x + np.sqrt(sigma2)*tr.randn(M*stacks,channel_uses).to(Device)\n",
    "\n",
    "def normalization(x): # E[|x|^2] = 1\n",
    "    return x / tr.sqrt((channel_uses*(x**2)).mean())\n",
    "\n",
    "def save():\n",
    "        tr.save({\n",
    "         'model_state_dict' : encoder.state_dict(), \n",
    "         'optimizer_state_dict': optimizer.state_dict(),\n",
    "         'loss': loss_history,\n",
    "         'constellations': Constellations,\n",
    "         'SNR': EsNo_dB,\n",
    "         'epochs' : epochs,\n",
    "         'stacks' : stacks,\n",
    "         'learning_rate': learning_rate,\n",
    "         'Device': Device,\n",
    "         'time': time.time()-start_time,\n",
    "         'Ldec': idx_train,\n",
    "         'Lbin': de2bi(idx_train,m)},'./Data/GMI/' + str(channel_uses) + 'D/' + str(M) + '/GMI_' + Estimation_type + '_' + str(channel_uses) + 'D_' + str(M) + '_' + str(EsNo_dB) + 'dB_'+ str(learning_rate)+'lr_' + Initial_Constellation)\n",
    "\n",
    "        sio.savemat('./Data/GMI/' + str(channel_uses) + 'D/' + str(M) + '/GMI_' + Estimation_type + '_' + str(channel_uses) + 'D_' + str(M) + '_' + str(EsNo_dB) + 'dB_'+ str(learning_rate)+'lr_' + Initial_Constellation + '.mat', {\n",
    "         'model_state_dict' : encoder.state_dict(), \n",
    "         'optimizer_state_dict': optimizer.state_dict(),\n",
    "         'loss': loss_history,\n",
    "         'constellations': Constellations,\n",
    "         'SNR': EsNo_dB,\n",
    "         'epochs' : epochs,\n",
    "         'stacks' : stacks,\n",
    "         'learning_rate': learning_rate,\n",
    "         'Device': Device,\n",
    "         'time': time.time()-start_time,\n",
    "         'Ldec': idx_train,\n",
    "         'Lbin': de2bi(idx_train,m)})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M = 16 #The cardinality size of the constellation\n",
    "stacks = 4 #The amount of noise samples per symbol for MC optimization\n",
    "channel_uses = 4 #The dimensionality of the constellation\n",
    "learning_rate = 0.1 #Learning Rate\n",
    "EsNo_dB = 9 # The SNR\n",
    "epochs = 800 #The amount of iterations\n",
    "Estimation_type = 'GH' #GH or MC\n",
    "Device = 'cpu' #Determines the device which the optimization is done on, 'cpu' for cpu and 'cuda:0', 'cuda:1' etc. for GPU\n",
    "Initial_Constellation = 'APSK' #You can choose either 'APSK', 'QAM', and in the case of 4D 'Random' or the directory to an optimized 2D constellation\n",
    "Amount_of_Rings  = 1 #Only relevant for APSK, tells you how many rings you want in the constellation, should be a power of 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Defining some of the parameters of the optimization\n",
    "m = int(np.log2(M)) #The amount of bits per symbol\n",
    "EsNo_r = 10**(EsNo_dB/10)\n",
    "sigma2 = 1/(channel_uses*EsNo_r) # noise variance per channel use\n",
    "GH =  sio.loadmat('GaussHermite_J_10.mat')#Loading in the Gauss-Hermite points\n",
    "X_eye = tr.eye(M).to(Device) # The input to our neural network\n",
    "X_tilde = X_eye.repeat(stacks,1).to(Device)\n",
    "if channel_uses == 2:\n",
    "    if Initial_Constellation == 'QAM':\n",
    "        [X_target,idx_train]= get_constellation(M)\n",
    "    else:\n",
    "        [X_target,idx_train] = get_APSK(M,Amount_of_Rings)\n",
    "        Initial_Constellation = Initial_Constellation + '_' + str(Amount_of_Rings) + 'Rings'\n",
    "else:\n",
    "    if Initial_Constellation == 'QAM':\n",
    "        [X_target,idx_train]= get_constellation_4D(M)\n",
    "    elif Initial_Constellation == 'APSK':\n",
    "        [X_target,idx_train] = get_APSK_4D(M,Amount_of_Rings)\n",
    "        Initial_Constellation = Initial_Constellation + '_' + str(Amount_of_Rings) + 'Rings'\n",
    "    elif Initial_Constellation == 'Random':\n",
    "        [X_target,idx_train] = get_Random_4D(M)\n",
    "    else:\n",
    "        Optimized_Constellation_2D = tr.load(Initial_Constellation)\n",
    "        X_2D = Optimized_Constellation_2D['constellations'][:,:,-1]\n",
    "        l_2D =  Optimized_Constellation_2D['Ldec']\n",
    "        [X_target,idx_train] = get_Optimized_4D(X_2D,l_2D,M)\n",
    "        Initial_Constellation = 'Optimized_2D'\n",
    "labeling = de2bi(np.arange(M), m)\n",
    "X_target = tr.tensor(X_target,dtype = tr.float32).to(Device)\n",
    "X_target = normalization(X_target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#=====================================================#\n",
    "# build the computation graph\n",
    "#=====================================================#\n",
    "def GMI(X_tilde, idx):\n",
    "    X = normalization(encoder(X_tilde))# minibatch_size x channel_uses\n",
    "    idx2 = np.zeros(M)\n",
    "    for i in range(M):\n",
    "        idx2[idx[i]] = i\n",
    "    X_sorted = tr.zeros((M*stacks,channel_uses)).to(Device)\n",
    "    X_temp = X[idx2,:]\n",
    "    for i in range(stacks):\n",
    "        X_sorted[i*M:(i+1)*M,:] = X_temp[:,:]\n",
    "    Y = awgn_channel(X_sorted)                 # minibatch_size x channel_uses\n",
    "    # compute posterior distribution\n",
    "    XX0_tmp,XX1_tmp,XX0,XX1 = [],[],[],[]\n",
    "    for i in range(int(m)):\n",
    "        XX0_tmp.clear()\n",
    "        XX1_tmp.clear()\n",
    "        for j in range(M):\n",
    "            if labeling[j,i] == 0:\n",
    "                XX0_tmp.append(tr.reshape(X_sorted[j,:], shape=[1,1,channel_uses]))\n",
    "            else:\n",
    "                XX1_tmp.append(tr.reshape(X_sorted[j,:], shape=[1,1,channel_uses]))\n",
    "        XX0.append(tr.cat(XX0_tmp, 0))\n",
    "        XX1.append(tr.cat(XX1_tmp, 0))\n",
    "    XB = []\n",
    "    for i in range(int(m)):\n",
    "        tmp = tr.cat([XX0[i].repeat([1, int(2**i), 1]), XX1[i].repeat([1, int(2**i), 1])], 1)\n",
    "        XB.append(tmp.repeat([1, int(2**(m-i-1)*stacks), 1]))\n",
    "    XX = tr.reshape(X_sorted[0:M:,:], shape=[M, 1, channel_uses])\n",
    "    YY = tr.reshape(Y, shape=[1, M*stacks, channel_uses])\n",
    "    epsilon = 1e-10 # to avoid log(x) = -Inf\n",
    "    loss = 0\n",
    "    for i in range(int(m)):\n",
    "        num = (tr.exp(-tr.sum((YY-XB[i])**2,2)/(2*sigma2))).mean(0) # f_{Y|B_i}\n",
    "        den = (tr.exp(-tr.sum((YY-XX)**2,2)/(2*sigma2))).mean(0) # f_{Y}\n",
    "        posterior = num/den\n",
    "        loss = loss - (tr.log(posterior + epsilon)/np.log(2)).mean()\n",
    "    return loss\n",
    "          \n",
    "          \n",
    "\n",
    "def GMI_GH(X_tilde,idx):\n",
    "    X = normalization(encoder(X_tilde))\n",
    "    idx2 = np.zeros(M)\n",
    "    for i in range(M):\n",
    "        idx2[idx[i]] = i \n",
    "    X = X[idx2,:]\n",
    "    Dmat = tr.zeros(M,M,channel_uses).to(Device)\n",
    "    Dmat[:,:,0] = X[:,0].unsqueeze(1) -(X[:,0].unsqueeze(1)).t() #Calculate the distances between constellation points\n",
    "    Dmat[:,:,1] = X[:,1].unsqueeze(1) -(X[:,1].unsqueeze(1)).t()\n",
    "    \n",
    "    Ik1 = np.zeros([int(M/2),int(m)]) #Find the pointers to the subconstellations\n",
    "    Ik0 = np.zeros([int(M/2),int(m)])\n",
    "    for kk in range(int(m)): \n",
    "        Ik1[:,kk] = np.where(labeling[:,kk] == 1)[0]\n",
    "        Ik0[:,kk] = np.where(labeling[:,kk] == 0)[0]\n",
    "    Ik1 = tr.tensor(Ik1, dtype = tr.int64).unsqueeze(1).to(Device)\n",
    "    Ikden1 = Ik1 + M*tr.transpose(Ik1,0,1)\n",
    "    Ik0 = tr.tensor(Ik0, dtype = tr.int64).unsqueeze(1).to(Device)\n",
    "    Ikden0 = Ik0 + M*tr.transpose(Ik0,0,1)\n",
    "    \n",
    "    GH_xi = tr.tensor(GH['xi'], dtype  = tr.float32).to(Device)#Load in the Gauss-Hermite points\n",
    "    GH_alpha = tr.tensor(GH['alpha'], dtype = tr.float32).to(Device)#Load in the Gauss-Hermite weigths\n",
    "    \n",
    "    Es = (X[:,0]**2 + X[:,1]**2).mean() #Calculate the signal energy\n",
    "    EsN0lin = 10**(EsNo_dB/10)  #Turn the SNR value from dB to a linear value\n",
    "    SigmaZ2 = (Es/(EsN0lin)) #Calculate the noise \n",
    "    \n",
    "    num = tr.zeros((M,M)).to(Device)\n",
    "    a = tr.zeros(M).to(Device)\n",
    "    num0 = tr.zeros((M,m)).to(Device)\n",
    "    num1 = tr.zeros((M,m)).to(Device)\n",
    "    den0 = tr.zeros((M,m)).to(Device)\n",
    "    den1 = tr.zeros((M,m)).to(Device)\n",
    "    sum_0 = tr.zeros((int(M/2),m)).to(Device)\n",
    "    sum_1 = tr.zeros((int(M/2),m)).to(Device)\n",
    "    l_2 = tr.log(tr.tensor(2, dtype = tr.float32)).to(Device)\n",
    "    GMI = 0\n",
    "    for l1 in range(10): #Dimension 1\n",
    "        for l2 in range(10): #Dimension 2\n",
    "             num = tr.exp(-((Dmat[:,:,0]**2 + Dmat[:,:,1]**2) + 2*tr.sqrt(SigmaZ2)*(GH_xi[l1]*Dmat[:,:,0] - GH_xi[l2]*Dmat[:,:,1]))/SigmaZ2)\n",
    "             a = num.sum(1)\n",
    "             num0 = tr.squeeze(a.take(Ik0))\n",
    "             den0 = tr.squeeze(num.take(Ikden0).sum(1))\n",
    "             den1 = tr.squeeze(num.take(Ikden1).sum(1))\n",
    "             num1 = tr.squeeze(a.take(Ik1))\n",
    "             sum_0 = GH_alpha[l1]*GH_alpha[l2]*tr.log(num0/den0)/l_2 +sum_0\n",
    "             sum_1 = GH_alpha[l1]*GH_alpha[l2]*tr.log(num1/den1)/l_2 +sum_1\n",
    "    sum_0 = tr.sum(sum_0)\n",
    "    sum_1 = tr.sum(sum_1)\n",
    "    GMI = m-1/M/math.pi*(sum_0 + sum_1)\n",
    "    loss = -GMI \n",
    "    return loss\n",
    "\n",
    "def GMI_GH_4D(X_tilde,idx):\n",
    "    X = normalization(encoder(X_tilde))\n",
    "    idx2 = np.zeros(M)\n",
    "    for i in range(M):\n",
    "        idx2[idx[i]] = i \n",
    "    X = X[idx2,:]\n",
    "    Dmat = tr.zeros(M,M,channel_uses).to(Device)\n",
    "    Dmat[:,:,0] = X[:,0].unsqueeze(1) -(X[:,0].unsqueeze(1)).t() #Calculate the distances between constellation points\n",
    "    Dmat[:,:,1] = X[:,1].unsqueeze(1) -(X[:,1].unsqueeze(1)).t()\n",
    "    Dmat[:,:,2] = X[:,2].unsqueeze(1) -(X[:,2].unsqueeze(1)).t()\n",
    "    Dmat[:,:,3] = X[:,3].unsqueeze(1) -(X[:,3].unsqueeze(1)).t()\n",
    "    \n",
    "    Ik1 = np.zeros([int(M/2),int(m)]) #Find the pointers to the subconstellations\n",
    "    Ik0 = np.zeros([int(M/2),int(m)])\n",
    "    for kk in range(int(m)): \n",
    "        Ik1[:,kk] = np.where(labeling[:,kk] == 1)[0]\n",
    "        Ik0[:,kk] = np.where(labeling[:,kk] == 0)[0]\n",
    "    Ik1 = tr.tensor(Ik1, dtype = tr.int64).unsqueeze(1).to(Device)\n",
    "    Ikden1 = Ik1 + M*tr.transpose(Ik1,0,1).to(Device)\n",
    "    Ik0 = tr.tensor(Ik0, dtype = tr.int64).unsqueeze(1).to(Device)\n",
    "    Ikden0 = Ik0 + M*tr.transpose(Ik0,0,1).to(Device)\n",
    "    \n",
    "    GH_xi = tr.tensor(GH['xi'], dtype  = tr.float32).to(Device)#Load in the Gauss-Hermite points\n",
    "    GH_alpha = tr.tensor(GH['alpha'], dtype = tr.float32).to(Device)#Load in the Gauss-Hermite weigths\n",
    "    \n",
    "    Es = (X[:,0]**2 + X[:,1]**2 + X[:,2]**2 + X[:,3]**2 ).mean() #Calculate the signal energy\n",
    "    EsN0lin = 10**(EsNo_dB/10)  #Turn the SNR value from dB to a linear value\n",
    "    SigmaZ2 = (Es/(EsN0lin)) #Calculate the noise \n",
    "    sum_0 = 0 #Initialize the sum \n",
    "    sum_1 = 0\n",
    "    Dmatnorm = (Dmat[:,:,0]**2 + Dmat[:,:,1]**2 + Dmat[:,:,2]**2 + Dmat[:,:,3]**2)\n",
    "    for l1 in range(10): #Dimension 1\n",
    "        for l2 in range(10): #Dimension 2\n",
    "            for l3 in range(10):\n",
    "                for l4 in range(10):\n",
    "                     num = tr.exp(-(Dmatnorm + tr.sqrt(2*SigmaZ2)*(GH_xi[l1]*Dmat[:,:,0] - GH_xi[l2]*Dmat[:,:,1] + GH_xi[l3]*Dmat[:,:,2] - GH_xi[l4]*Dmat[:,:,3]))/(0.5*SigmaZ2))\n",
    "                     a = num.sum(1)                  \n",
    "                     num0 = tr.squeeze(a.take(Ik0))\n",
    "                     den0 = tr.squeeze(num.take(Ikden0).sum(1))\n",
    "                     den1 = tr.squeeze(num.take(Ikden1).sum(1))\n",
    "                     num1 = tr.squeeze(a.take(Ik1))\n",
    "                     sum_0 = GH_alpha[l1]*GH_alpha[l2]*GH_alpha[l3]*GH_alpha[l4]*tr.log(num0/den0)/tr.log(tr.tensor(2, dtype = tr.float32))  + sum_0   \n",
    "                     sum_1 = GH_alpha[l1]*GH_alpha[l2]*GH_alpha[l3]*GH_alpha[l4]*tr.log(num1/den1)/tr.log(tr.tensor(2, dtype = tr.float32))  + sum_1\n",
    "    sum_0 = tr.sum(sum_0)\n",
    "    sum_1 = tr.sum(sum_1)\n",
    "    GMI = m-1/M/((math.pi)**2)*(sum_0 + sum_1)\n",
    "    loss = -GMI \n",
    "    return loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Training the model\n",
    "start_time = time.time()\n",
    "loss_history = np.zeros(epochs) # For saving the losses\n",
    "Constellations = np.zeros((M,channel_uses,epochs)) #For saving the constellations\n",
    "e = 0\n",
    "encoder = tr.nn.Sequential()\n",
    "encoder.add_module('last', tr.nn.Linear(M,channel_uses,bias = False))\n",
    "encoder.to(Device)\n",
    "optimizer = tr.optim.Adam(encoder.parameters(), learning_rate)\n",
    "while e == 0:\n",
    "    loss = sum(sum(((X_target - normalization(encoder(X_eye)))**2)))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if loss < 10**-8:\n",
    "        e = 1\n",
    "for i in range(1, epochs+1):\n",
    "    if Estimation_type == 'GH':# and i > 500: #Uncomment this if you want the first several iterations to be calculated using MC, advised when doing 4D optimizations\n",
    "        if channel_uses ==2:\n",
    "            loss = GMI_GH(X_eye, idx_train)\n",
    "        else:\n",
    "            loss = GMI_GH_4D(X_eye, idx_train)\n",
    "    else:\n",
    "        loss = GMI(X_tilde, idx_train)\n",
    "    Constellations[:,:,i-1] = normalization(encoder(X_eye)).detach().cpu().numpy()\n",
    "    if i%100 == 0 or i ==1:\n",
    "        print('iter ', i, ' loss', loss, 'time', time.time()-start_time)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    loss_history[i-1] = loss.detach().cpu().numpy()\n",
    "    save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Plotting the final constellation with labeling\n",
    "constellation = Constellations[:,:,-1]\n",
    "x = constellation[:, 0] #Change from 0 to 2 for the second 2D for 4D constellations\n",
    "y = constellation[:, 1] #Change from 1 to 3 for the second 2D for 4D constellations\n",
    "max_axis = 1.8\n",
    "fig = plt.figure(figsize=(m*3-3,3))\n",
    "\n",
    "for i in range(int(m)):\n",
    "\tax1 = plt.subplot(1, m, i+1)\n",
    "\tax1.scatter(x[labeling[:,i] == 0], y[labeling[:,i] == 0], c='b', marker='.')\n",
    "\tax1.scatter(x[labeling[:,i] == 1], y[labeling[:,i] == 1], c='r', marker='.')\n",
    "\t\n",
    "\tplt.xlabel('X')\n",
    "\tplt.ylabel('Y')\n",
    "\tplt.grid()\n",
    "\t#plt.axis('equal')\n",
    "\tplt.gca().set_aspect('equal', adjustable='box')\n",
    "\tplt.xlim(-max_axis, max_axis)\n",
    "\tplt.ylim(-max_axis, max_axis)\n",
    "\tax1.set_title('bit position {}'.format(i))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Creating an animation of the evolution of the constellation, only works for 2D constellations\n",
    "plt.rcParams['animation.embed_limit'] = 2**128\n",
    "def init():\n",
    "    line[0].set_data([],[])\n",
    "    line[1].set_data([],[])\n",
    "    return(line,)\n",
    "\n",
    "def animate(i):\n",
    "    x = Constellations[:,0,i]\n",
    "    y = Constellations[:,1,i]\n",
    "    line[0].set_data(x,y)\n",
    "    line[1].set_data(np.linspace(0,i,i),loss_history[0:i])\n",
    "    return(line,)\n",
    "\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(1,2,1)\n",
    "ax2 = fig.add_subplot(1,2,2)\n",
    "ax.set_xlim((-1.8,1.8))\n",
    "ax.set_ylim((-1.8,1.8))\n",
    "ax2.set_xlim((-1,epochs))\n",
    "ax2.set_ylim((-8.75,-8.6))\n",
    "ax.set_aspect('equal')\n",
    "ax2.set_xlabel('iterations')\n",
    "ax2.set_ylabel('loss')\n",
    "line1, = ax.plot([],[],'.',lw=2,)\n",
    "line2, = ax2.plot([],[],lw=2,)\n",
    "line1.set_color('b')\n",
    "line2.set_color('b')\n",
    "line = [line1, line2]\n",
    "anim = animation.FuncAnimation(fig,animate,init_func = init, frames = epochs, interval = 10, blit = False)\n",
    "anim.save('./Data/GMI/' + str(channel_uses) + 'D/' + str(M) + '/GMI_' + Estimation_type + '_' + str(channel_uses) + 'D_' + str(M) + '_' + str(EsNo_dB) + 'dB_'+ str(learning_rate)+'lr_' + Initial_Constellation + '_animation.mp4') \n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
